package org.example.game.weekly.game410;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

/**
 * @Description: TODO
 * @Author wyatt
 * @Data 2024/08/03 20:46
 */
public class Solution002 {

    public static void main(String[] args) {
        System.out.println(new Solution002().countGoodNodes(new int[][]{{0, 1}, {1, 2}, {2, 3}, {3, 4}, {0, 5}, {1, 6}, {2, 7}, {3, 8}}));
        System.out.println(new Solution002().countGoodNodes(new int[][]{{0,1},{1,2},{1,3},{1,4},{0,5},{5,6},{6,7},{7,8},{0,9},{9,10},{9,12},{10,11}}));
        System.out.println(new Solution002().countGoodNodes(new int[][]{{6,0},{1,0},{5,1},{2,5},{3,1},{4,3}}));
    }


    public int countGoodNodes(int[][] edges) {
        // 1. N个节点
        // 2. N个节点的子节点 的数量相同， 则改节点为好节点   或者只有0个或者1个子节点
        // 3. 节点的数量 = 下一个层级的数量，
        HashMap<Integer, List<Integer>> treeMap = new HashMap<>();

        for (int i = 0; i < edges.length; i++) {
            treeMap.putIfAbsent(edges[i][0], new ArrayList<>());
            treeMap.putIfAbsent(edges[i][1], new ArrayList<>());
            treeMap.get(edges[i][0]).add(edges[i][1]);
        }

        int sum = 0;
        for (Integer node : treeMap.keySet()) {
            List<Integer> nextNodeList = treeMap.get(node);
            if(nextNodeList.size() <= 1){
                sum++;
            }else {
                //判断子树的节点数量
                int count = -1;
                boolean match = true;
                for (Integer nextNode : nextNodeList) {
                    int nextCount = getLenOfNode(treeMap, nextNode);
                    if(count == -1){
                        count = nextCount;
                    }else if(count != nextCount){
                        match = false;
                        break;
                    }
                }
                if(match){
                    sum++;
                }
            }
        }
        return sum;
    }

    private int getLenOfNode(HashMap<Integer, List<Integer>> treeMap, Integer nextNode) {
        if(!treeMap.containsKey(nextNode)){
            return 0;
        }
        int sum = 0;
        sum+=treeMap.get(nextNode).size();
        for(int node : treeMap.get(nextNode)){
            sum += getLenOfNode(treeMap, node);
        }
        return sum;
    }

}
